# -*- coding: utf-8 -*-
# data_processing1.py
from turtle import st

import pandas as pd
import re
from zhipuai import ZhipuAI
from typing import Dict, List, Union

# -------------------- 全局配置 --------------------
DATA_FILE = "拉勾网2023招聘数据.csv"


# -------------------- 数据预处理 --------------------
def load_data() -> pd.DataFrame:
    """数据加载和清洗函数"""

    def parse_salary(s):
        try:
            if '薪' in s:
                s = s.split('·')[0]
            parts = s.lower().replace('k', '-').split('-')
            nums = [float(n) * 1000 for n in parts if n.replace('.', '').isdigit()]
            return sum(nums) / len(nums) if nums else 0
        except:
            return 0

    df = pd.read_csv(DATA_FILE)
    df.columns = df.columns.str.lower().str.replace(' ', '')

    # 列名标准化
    df.rename(columns={
        'positionname': 'position',
        'workyear': 'workyear',
        'companyname': 'company'
    }, inplace=True, errors='ignore')

    # 数据清洗流程
    required_columns = {'city', 'position', 'salary', 'education', 'workyear'}
    if missing := required_columns - set(df.columns):
        raise ValueError(f"缺少必要字段: {missing}")

    df = df.dropna(subset=['city', 'position', 'salary'])
    df['position'] = df['position'].str.lower()
    df['city'] = df['city'].str.lower()
    df['salary'] = df['salary'].apply(parse_salary)

    # 公司信息处理
    company_columns = ['company', 'companyname', '公司名称']
    for col in company_columns:
        if col in df.columns:
            df['company'] = df[col].str.replace(r'[（(].*?[）)]', '', regex=True)
            break
    else:
        df['company'] = '未知'

    # 工作经验处理
    df['workyear'] = df['workyear'].astype(str).apply(
        lambda x: x.replace("经验", "").replace("年", "").replace("以上", "+").strip()
    )

    # 学历处理
    df['education'] = df['education'].str.lower().str.replace(' ', '')

    return df


# -------------------- 业务逻辑 --------------------
class AnalysisEngine:
    @staticmethod
    def count_positions(df: pd.DataFrame, city: str, position: str) -> int:
        filtered = df[(df['city'] == city.lower()) & (df['position'] == position.lower())]
        return len(filtered)

    @staticmethod
    def average_salary(df: pd.DataFrame, city: str, position: str) -> str:
        filtered = df[(df['city'] == city.lower()) & (df['position'] == position.lower())]
        return f"{filtered['salary'].mean() / 1000:.1f}k" if not filtered.empty else "无数据"

    @staticmethod
    def compare_salaries(df: pd.DataFrame, city1: str, city2: str, position: str) -> str:
        avg1 = df[(df['city'] == city1.lower()) & (df['position'] == position.lower())]['salary'].mean() or 0
        avg2 = df[(df['city'] == city2.lower()) & (df['position'] == position.lower())]['salary'].mean() or 0
        return f"{city1}: {avg1 / 1000:.1f}k vs {city2}: {avg2 / 1000:.1f}k"


# -------------------- AI 交互核心 --------------------
class AIClient:
    def __init__(self, df: pd.DataFrame):
        self.client = ZhipuAI(api_key=st.session_state.api_key)
        self.df = df
        self._init_tools()

    def _init_tools(self):
        """初始化函数工具集"""
        self.tools = [
            {
                "type": "function",
                "function": {
                    "name": "get_position_stats",
                    "description": "获取职位统计数据",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "city": {"type": "string"},
                            "position": {"type": "string"},
                            "metric": {
                                "type": "string",
                                "enum": ["count", "avg_salary", "range"]
                            }
                        },
                        "required": ["city", "position", "metric"]
                    }
                }
            }
        ]

    def process_query(self, query: str) -> Dict:
        """处理自然语言查询"""
        response = self.client.chat.completions.create(
            model="glm-4",
            messages=[{"role": "user", "content": query}],
            tools=self.tools
        )
        return response.choices[0].message.tool_calls[0].function

    def execute_query(self, func_name: str, params: Dict) -> Union[str, Dict]:
        """执行查询操作"""
        if func_name == "get_position_stats":
            city = params["city"].lower()
            position = params["position"].lower()
            metric = params["metric"]

            filtered = self.df[
                (self.df['city'] == city) &
                (self.df['position'] == position)
                ]

            if metric == "count":
                return {"count": len(filtered)}
            elif metric == "avg_salary":
                return {"average_salary": filtered['salary'].mean()}
            elif metric == "range":
                return {
                    "min": filtered['salary'].min(),
                    "max": filtered['salary'].max()
                }
        return {"error": "未知查询类型"}